{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "23913a93",
   "metadata": {},
   "source": [
    "注意需要在本地安装FFmpeg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ee01b11-c385-4803-bba6-77596ea76393",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from datasets import load_dataset\n",
    "from transformers import WhisperProcessor, WhisperForConditionalGeneration\n",
    "from moviepy.editor import VideoFileClip\n",
    "import collections.abc\n",
    "import ffmpeg\n",
    "import librosa\n",
    "import soundfile as sf\n",
    "\n",
    "import os\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0972f845",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_flac_from_mp4(mp4_file, flac_file):\n",
    "    \"\"\"\n",
    "    Extracts audio from an MP4 file and saves it as FLAC.\n",
    "\n",
    "    Args:\n",
    "        mp4_file: Path to the MP4 file.\n",
    "        flac_file: Path to the output FLAC file.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        input_stream = ffmpeg.input(mp4_file)\n",
    "        print(f'Input {mp4_file} successfully')\n",
    "        \n",
    "        audio_stream = input_stream.audio\n",
    "        print(f'Get audio from {mp4_file} successfully')\n",
    "        \n",
    "        output_stream = ffmpeg.output(audio_stream, flac_file, acodec='flac')\n",
    "        print('Convert audio into flac successfully')        \n",
    "        \n",
    "        ffmpeg.run(output_stream)\n",
    "        print(f\"Successfully extracted audio to {flac_file}\")\n",
    "    except ffmpeg.Error as e:\n",
    "        print(f\"Error extracting audio: {e}\")\n",
    "    \n",
    "def load_audio_from_flac(file_path):\n",
    "    \"\"\"Loads audio data from a FLAC file.\n",
    "\n",
    "    Args:\n",
    "    file_path (str): Path to the FLAC file.\n",
    "\n",
    "    Returns:\n",
    "    tuple: A tuple containing the audio data as a NumPy array and the sample rate.\n",
    "    \"\"\"\n",
    "    sample_rate_whisper = 16000\n",
    "    \n",
    "    try:\n",
    "        audio_data, sample_rate = librosa.load(file_path, sr=sample_rate_whisper)\n",
    "        return audio_data, sample_rate\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading audio: {e}\")\n",
    "        return None, None\n",
    "    \n",
    "# Define a function to process audio in chunks\n",
    "def transcribe_long_audio(model, audio_array, sampling_rate):\n",
    "    chunk_length = 30  # Process in 30-second chunks\n",
    "    num_chunks = int(len(audio_array) / (chunk_length * sampling_rate)) + 1\n",
    "    transcriptions = []\n",
    "\n",
    "    for i in range(num_chunks):\n",
    "        start = i * chunk_length * sampling_rate\n",
    "        end = min((i + 1) * chunk_length * sampling_rate, len(audio_array))\n",
    "        chunk = audio_array[start:end]\n",
    "\n",
    "        # Use the model and processor to transcribe the audio:\n",
    "        input_features = processor(\n",
    "            chunk,\n",
    "            sampling_rate=sampling_rate,\n",
    "            return_tensors=\"pt\",\n",
    "            return_attention_mask=True\n",
    "        )\n",
    "\n",
    "        # Generate token ids\n",
    "        predicted_ids = model.generate(\n",
    "            input_features=input_features.input_features,\n",
    "            attention_mask=input_features.attention_mask\n",
    "        )\n",
    "\n",
    "        # Decode token ids to text\n",
    "        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
    "\n",
    "        # Joint outputs of chunks\n",
    "        transcriptions.extend(transcription)\n",
    "\n",
    "    return \" \".join(transcriptions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a39775cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "video_file = \".\\Data\\President20Kennedy.mp4\"\n",
    "audio_file = \".\\Data\\President20Kennedy.flac\"\n",
    "\n",
    "if not os.path.exists(audio_file):\n",
    "    extract_flac_from_mp4(video_file, audio_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22f02dd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_data, sample_rate = load_audio_from_flac(audio_file)\n",
    "\n",
    "if audio_data is not None:\n",
    "    print(f\"Audio loaded successfully. Sample rate: {sample_rate} Hz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "21d87405",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the Whisper model in Hugging Face format:\n",
    "processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny.en\")\n",
    "model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-tiny.en\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "857e6695",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the model and processor to transcribe the audio:\n",
    "input_features = processor(\n",
    "    audio_data,\n",
    "    sampling_rate=sample_rate,\n",
    "    return_tensors=\"pt\",\n",
    "    return_attention_mask=True\n",
    ")\n",
    "\n",
    "# Generate token ids\n",
    "predicted_ids = model.generate(\n",
    "    input_features=input_features.input_features,\n",
    "    attention_mask=input_features.attention_mask\n",
    ")\n",
    "\n",
    "# Decode token ids to text\n",
    "transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
    "\n",
    "print(transcription)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb53a714",
   "metadata": {},
   "outputs": [],
   "source": [
    "transcription = transcribe_long_audio(model, audio_data, sample_rate)\n",
    "\n",
    "print(transcription)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31e7d49f",
   "metadata": {},
   "source": [
    "Connect to local LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "0872b337",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_url = 'http://localhost:8080'\n",
    "\n",
    "def get_server_health():\n",
    "    global base_url\n",
    "    \n",
    "    response = requests.get(f'{base_url}/health')\n",
    "    return response.json()\n",
    "\n",
    "def post_completion(context, user_input):\n",
    "    global base_url\n",
    "    \n",
    "    prompt = f\"{context}\\nUser: {user_input}\\nAssistant:\"\n",
    "    data = {\n",
    "        'prompt': prompt,\n",
    "        'temperature': 0.5,\n",
    "        'top_k': 35,\n",
    "        'top_p': 0.95,\n",
    "        'n_predict': [],\n",
    "        'stop': [\"</s>\", \"Assistant:\", \"User:\"]\n",
    "    }\n",
    "    headers = {'Content-Type': 'application/json'}\n",
    "    response = requests.post(f'{base_url}/completion', json=data, headers=headers)\n",
    "    if response.status_code == 200:\n",
    "        return response.json()['content'].strip()\n",
    "    else:\n",
    "        return \"Error processing your request. Please try again.\"\n",
    "    \n",
    "def update_context(context, user_input, assistant_response):\n",
    "    return f\"{context}\\nUser: {user_input}\\nAssistant: {assistant_response}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39d21f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "health = get_server_health()\n",
    "print('Server Health:', health)\n",
    "\n",
    "text = transcription\n",
    "user_input = 'Please read the following article and provide a concise summary of the main points and key details of what the speaker talked. The outputs should include the main point and at least three detailed of the main points. The summary should be no longer than 200 words.\\n'\n",
    "\n",
    "print('User:', user_input)\n",
    "assistant_response = post_completion(user_input, text)\n",
    "print('Assistant:', assistant_response)\n",
    "\n",
    "# if health.get('status') == 'ok':\n",
    "#     while True:\n",
    "#         user_input = input(\"Enter a prompt or type 'exit' to quit: \")\n",
    "#         if user_input.lower() == 'exit':\n",
    "#             break\n",
    "        \n",
    "#         print('User:', user_input)\n",
    "#         assistant_response = post_completion(context, user_input)\n",
    "#         print('Assistant:', assistant_response)\n",
    "\n",
    "#         context = update_context(context, user_input, assistant_response)\n",
    "# else:\n",
    "#     print(\"Server is not ready for requests.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
